msg_tool\scripts\artemis/
asb.rs

1//! Artemis Engine ASB file (.asb/.iet)
2use crate::ext::io::*;
3use crate::scripts::base::*;
4use crate::types::*;
5use crate::utils::encoding::*;
6use crate::utils::escape::*;
7use anyhow::Result;
8use serde::{Deserialize, Serialize};
9use std::collections::BTreeMap;
10use std::io::{Read, Write};
11use std::ops::Index;
12use stylua_lib::{Config as LuaFormatterConfig, OutputVerification, format_code};
13use unicode_segmentation::UnicodeSegmentation;
14
15#[derive(Debug)]
16/// The builder for Artemis ASB scripts.
17pub struct ArtemisAsbBuilder {}
18
19impl ArtemisAsbBuilder {
20    /// Creates a new instance of `ArtemisAsbBuilder`.
21    pub fn new() -> Self {
22        ArtemisAsbBuilder {}
23    }
24}
25
26impl ScriptBuilder for ArtemisAsbBuilder {
27    fn default_encoding(&self) -> Encoding {
28        Encoding::Utf8
29    }
30
31    fn build_script(
32        &self,
33        buf: Vec<u8>,
34        filename: &str,
35        encoding: Encoding,
36        _archive_encoding: Encoding,
37        config: &ExtraConfig,
38        _archive: Option<&Box<dyn Script>>,
39    ) -> Result<Box<dyn Script>> {
40        Ok(Box::new(Asb::new(buf, encoding, config, filename)?))
41    }
42
43    fn extensions(&self) -> &'static [&'static str] {
44        &["asb", "iet"]
45    }
46
47    fn script_type(&self) -> &'static ScriptType {
48        &ScriptType::ArtemisAsb
49    }
50
51    fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
52        if buf_len >= 5 && buf.starts_with(b"ASB\0\0") {
53            return Some(20);
54        }
55        None
56    }
57
58    fn can_create_file(&self) -> bool {
59        true
60    }
61
62    fn create_file<'a>(
63        &'a self,
64        filename: &'a str,
65        writer: Box<dyn WriteSeek + 'a>,
66        encoding: Encoding,
67        file_encoding: Encoding,
68        config: &ExtraConfig,
69    ) -> Result<()> {
70        create_file(
71            filename,
72            writer,
73            encoding,
74            file_encoding,
75            config.custom_yaml,
76        )
77    }
78}
79
80fn escape_text(s: &str) -> String {
81    let mut escaped = String::with_capacity(s.len());
82    for c in s.chars() {
83        match c {
84            '&' => escaped.push_str("&amp;"),
85            '<' => escaped.push_str("&lt;"),
86            _ => escaped.push(c),
87        }
88    }
89    escaped
90}
91
92#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
93struct Command {
94    pub name: String,
95    pub line_number: u32,
96    pub attributes: BTreeMap<String, String>,
97}
98
99impl Command {
100    pub fn new(name: String, line_number: u32) -> Self {
101        Command {
102            name,
103            line_number,
104            attributes: BTreeMap::new(),
105        }
106    }
107
108    pub fn to_xml(&self) -> String {
109        let mut xml = format!("<{}", self.name);
110        for (key, value) in &self.attributes {
111            xml.push_str(&format!(" {}=\"{}\"", key, escape_xml_text_value(value)));
112        }
113        xml.push('>');
114        xml
115    }
116}
117
118impl<'a> Index<&'a str> for Command {
119    type Output = str;
120    fn index(&self, key: &'a str) -> &Self::Output {
121        self.attributes.get(key).map_or("", |s| s.as_str())
122    }
123}
124
125impl<'a> Index<&'a String> for Command {
126    type Output = str;
127    fn index(&self, key: &'a String) -> &Self::Output {
128        self.attributes.get(key).map_or("", |s| s.as_str())
129    }
130}
131
132impl Index<String> for Command {
133    type Output = str;
134    fn index(&self, key: String) -> &Self::Output {
135        self.attributes.get(&key).map_or("", |s| s.as_str())
136    }
137}
138
139#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
140#[serde(untagged)]
141enum Item {
142    Command(Command),
143    Label(String),
144}
145
146impl Item {
147    pub fn is_command(&self) -> bool {
148        matches!(self, Item::Command(_))
149    }
150
151    pub fn is_command_name(&self, name: &str) -> bool {
152        if let Item::Command(cmd) = self {
153            cmd.name == name
154        } else {
155            false
156        }
157    }
158}
159
160trait CustomReadFn {
161    fn read_string(&mut self, encoding: Encoding) -> Result<String>;
162    fn read_item(&mut self, encoding: Encoding) -> Result<Item>;
163}
164
165impl<T: Read> CustomReadFn for T {
166    fn read_string(&mut self, encoding: Encoding) -> Result<String> {
167        let len = self.read_u32()?;
168        let data = self.read_exact_vec(len as usize)?;
169        if self.read_u8()? != 0 {
170            return Err(anyhow::anyhow!("String not null-terminated"));
171        }
172        let s = decode_to_string(encoding, &data, true)?;
173        Ok(s)
174    }
175
176    fn read_item(&mut self, encoding: Encoding) -> Result<Item> {
177        let typ = self.read_u32()?;
178        match typ {
179            0 => {
180                let name = self.read_string(encoding)?;
181                let line_number = self.read_u32()?;
182                let mut command = Command::new(name, line_number);
183                let attr_count = self.read_u32()?;
184                for _ in 0..attr_count {
185                    let key = self.read_string(encoding)?;
186                    let value = self.read_string(encoding)?;
187                    command.attributes.insert(key, value);
188                }
189                Ok(Item::Command(command))
190            }
191            1 => {
192                let label = self.read_string(encoding)?;
193                Ok(Item::Label(label))
194            }
195            _ => {
196                return Err(anyhow::anyhow!("Unknown item type: {}", typ));
197            }
198        }
199    }
200}
201
202trait CustomWriteFn {
203    fn write_string(&mut self, s: &str, encoding: Encoding) -> Result<()>;
204    fn write_item(&mut self, item: &Item, encoding: Encoding) -> Result<()>;
205}
206
207impl<T: Write> CustomWriteFn for T {
208    fn write_string(&mut self, s: &str, encoding: Encoding) -> Result<()> {
209        let data = encode_string(encoding, s, false)?;
210        self.write_u32(data.len() as u32)?;
211        self.write_all(&data)?;
212        self.write_u8(0)?; // Null-terminated
213        Ok(())
214    }
215
216    fn write_item(&mut self, item: &Item, encoding: Encoding) -> Result<()> {
217        match item {
218            Item::Command(cmd) => {
219                self.write_u32(0)?; // Type 0 for Command
220                self.write_string(&cmd.name, encoding)?;
221                self.write_u32(cmd.line_number)?;
222                self.write_u32(cmd.attributes.len() as u32)?;
223                for (key, value) in &cmd.attributes {
224                    self.write_string(key, encoding)?;
225                    self.write_string(value, encoding)?;
226                }
227            }
228            Item::Label(label) => {
229                self.write_u32(1)?; // Type 1 for Label
230                self.write_string(label, encoding)?;
231            }
232        }
233        Ok(())
234    }
235}
236
237struct TextParser<'a> {
238    items: Vec<Item>,
239    text: Vec<&'a str>,
240    pos: usize,
241    len: usize,
242    hcls_index: usize,
243}
244
245impl<'a> TextParser<'a> {
246    pub fn new(str: &'a str, hcls_index: usize) -> Self {
247        let text: Vec<&'a str> = UnicodeSegmentation::graphemes(str, true).collect();
248        let len = text.len();
249        TextParser {
250            items: Vec::new(),
251            text,
252            pos: 0,
253            len,
254            hcls_index,
255        }
256    }
257
258    pub fn parse(mut self) -> Result<Vec<Item>> {
259        while let Some(c) = self.peek() {
260            match c {
261                "<" => {
262                    self.parse_tag()?;
263                }
264                _ => {
265                    let mut text = String::new();
266                    self.eat_char();
267                    text.push_str(c);
268                    while let Some(b) = self.peek() {
269                        if b == "<" {
270                            break;
271                        }
272                        text.push_str(b);
273                        self.eat_char();
274                    }
275                    if !text.is_empty() {
276                        self.items.push(Item::Command(Command {
277                            name: "print".to_string(),
278                            line_number: 0,
279                            attributes: [("data".to_string(), unescape_xml(&text))].into(),
280                        }))
281                    }
282                }
283            }
284        }
285        let mut hcls = Command::new("hcls".to_string(), 0);
286        hcls.attributes
287            .insert("0".to_string(), self.hcls_index.to_string());
288        self.items.push(Item::Command(hcls));
289        Ok(self.items)
290    }
291
292    fn parse_tag(&mut self) -> Result<()> {
293        self.parse_indent("<")?;
294        let key = self.parse_key()?;
295        self.erase_whitespace();
296        let mut cmd = Command::new(key, 0);
297        loop {
298            let c = self.peek().ok_or(self.error2("Unexpected eof"))?;
299            match c {
300                ">" => {
301                    self.eat_char();
302                    break;
303                }
304                " " => {
305                    self.eat_char();
306                    continue;
307                }
308                _ => {
309                    let key = self.parse_key()?;
310                    self.parse_indent("=")?;
311                    let value = self.parse_str()?;
312                    cmd.attributes.insert(key, value);
313                }
314            }
315        }
316        self.items.push(Item::Command(cmd));
317        Ok(())
318    }
319
320    fn parse_key(&mut self) -> Result<String> {
321        self.erase_whitespace();
322        let mut key = String::new();
323        while let Some(c) = self.peek() {
324            if c == "=" || c == " " || c == ">" {
325                break;
326            }
327            key.push_str(c);
328            self.eat_char();
329        }
330        if key.is_empty() {
331            return self.error("Expected key, but found nothing");
332        }
333        Ok(key)
334    }
335
336    fn parse_str(&mut self) -> Result<String> {
337        self.erase_whitespace();
338        self.parse_indent("\"")?;
339        let mut text = String::new();
340        loop {
341            match self.next().ok_or(self.error2("Unexpected eof"))? {
342                "\"" => {
343                    break;
344                }
345                t => {
346                    text.push_str(t);
347                }
348            }
349        }
350        Ok(unescape_xml(&text))
351    }
352
353    fn erase_whitespace(&mut self) {
354        while let Some(c) = self.peek() {
355            if c == " " {
356                self.eat_char();
357            } else {
358                break;
359            }
360        }
361    }
362
363    fn parse_indent(&mut self, indent: &str) -> Result<()> {
364        for ident in indent.graphemes(true) {
365            match self.next() {
366                Some(c) => {
367                    if c != ident {
368                        return self.error("Unexpected indent");
369                    }
370                }
371                None => return self.error("Unexpected eof"),
372            }
373        }
374        Ok(())
375    }
376
377    fn eat_char(&mut self) {
378        if self.pos < self.len {
379            self.pos += 1;
380        }
381    }
382
383    fn next(&mut self) -> Option<&'a str> {
384        if self.pos < self.len {
385            let item = self.text[self.pos];
386            self.pos += 1;
387            Some(item)
388        } else {
389            None
390        }
391    }
392
393    fn peek(&self) -> Option<&'a str> {
394        if self.pos < self.len {
395            Some(self.text[self.pos])
396        } else {
397            None
398        }
399    }
400
401    fn error2<T>(&self, msg: T) -> anyhow::Error
402    where
403        T: std::fmt::Display,
404    {
405        anyhow::anyhow!("Failed to parse at position {}: {}", self.pos, msg)
406    }
407
408    fn error<T, A>(&self, msg: T) -> Result<A>
409    where
410        T: std::fmt::Display,
411    {
412        Err(anyhow::anyhow!(
413            "Failed to parse at position {}: {}",
414            self.pos,
415            msg
416        ))
417    }
418}
419
420#[derive(Debug)]
421/// The Artemis ASB script.
422pub struct Asb {
423    items: Vec<Item>,
424    custom_yaml: bool,
425    is_iet: bool,
426    format_lua: bool,
427}
428
429impl Asb {
430    /// Creates a new Artemis ASB script from the given buffer.
431    ///
432    /// * `buf` - The buffer containing the ASB data.
433    /// * `encoding` - The encoding used for the ASB data.
434    /// * `config` - Extra configuration options.
435    pub fn new(
436        buf: Vec<u8>,
437        encoding: Encoding,
438        config: &ExtraConfig,
439        filename: &str,
440    ) -> Result<Self> {
441        let mut data = MemReader::new(buf);
442        let mut magic = [0; 5];
443        data.read_exact(&mut magic)?;
444        if &magic != b"ASB\0\0" {
445            return Err(anyhow::anyhow!("Invalid ASB magic number: {:?}", magic));
446        }
447        let nums = data.read_u32()?;
448        let mut items = Vec::with_capacity(nums as usize);
449        for _ in 0..nums {
450            items.push(data.read_item(encoding)?);
451        }
452        Ok(Asb {
453            items,
454            custom_yaml: config.custom_yaml,
455            is_iet: std::path::Path::new(filename)
456                .extension()
457                .map_or(false, |ext| ext.eq_ignore_ascii_case("iet")),
458            format_lua: config.artemis_asb_format_lua,
459        })
460    }
461
462    fn to_string(&self, items: &[Item]) -> Result<String> {
463        if self.custom_yaml {
464            Ok(serde_yaml_ng::to_string(items)?)
465        } else {
466            Ok(serde_json::to_string_pretty(items)?)
467        }
468    }
469
470    fn format_lua(&self, script: &str) -> Result<String> {
471        let mut config = LuaFormatterConfig::new();
472        config.indent_type = stylua_lib::IndentType::Spaces;
473        config.indent_width = 2;
474        config.column_width = 120;
475        config.line_endings = stylua_lib::LineEndings::Unix;
476        Ok(format_code(script, config, None, OutputVerification::None)?)
477    }
478}
479
480impl Script for Asb {
481    fn default_output_script_type(&self) -> OutputScriptType {
482        if self.is_iet {
483            OutputScriptType::Custom
484        } else {
485            OutputScriptType::Json
486        }
487    }
488
489    fn is_output_supported(&self, out: OutputScriptType) -> bool {
490        if self.is_iet {
491            matches!(out, OutputScriptType::Custom)
492        } else {
493            true
494        }
495    }
496
497    fn default_format_type(&self) -> FormatOptions {
498        FormatOptions::None
499    }
500
501    fn custom_output_extension<'a>(&'a self) -> &'a str {
502        if self.custom_yaml { "yaml" } else { "json" }
503    }
504
505    fn extract_messages(&self) -> Result<Vec<Message>> {
506        let mut messages = Vec::new();
507        let mut name = None;
508        let mut cur_mes = String::new();
509        let mut in_print = false;
510        for item in self.items.iter() {
511            if in_print {
512                if let Item::Command(cmd) = item {
513                    match cmd.name.as_str() {
514                        "hcls" => {
515                            in_print = false;
516                            messages.push(Message {
517                                name: name.take(),
518                                message: cur_mes,
519                            });
520                            cur_mes = String::new();
521                        }
522                        "print" => {
523                            cur_mes.push_str(&escape_text(&cmd["data"]));
524                        }
525                        "rt" => {
526                            cur_mes.push('\n');
527                        }
528                        _ => {
529                            cur_mes.push_str(&cmd.to_xml());
530                        }
531                    }
532                    continue;
533                }
534            }
535            if let Item::Command(cmd) = item {
536                match cmd.name.as_str() {
537                    "print" => {
538                        cur_mes.push_str(&escape_text(&cmd["data"]));
539                        in_print = true;
540                    }
541                    "name" => {
542                        let v = (cmd.attributes.len() - 1).to_string();
543                        name = Some(cmd[v].to_owned());
544                    }
545                    "sel_text" => {
546                        let t = &cmd["text"];
547                        if !t.is_empty() {
548                            messages.push(Message {
549                                name: None,
550                                message: t.to_owned(),
551                            });
552                        }
553                    }
554                    "RegisterTextToHistory" => {
555                        let t = &cmd["1"];
556                        if !t.is_empty() {
557                            messages.push(Message {
558                                name: None,
559                                message: t.to_owned(),
560                            });
561                        }
562                    }
563                    _ => {}
564                }
565            }
566        }
567        if !cur_mes.is_empty() {
568            messages.push(Message {
569                name: name.take(),
570                message: cur_mes,
571            });
572        }
573        Ok(messages)
574    }
575
576    fn import_messages<'a>(
577        &'a self,
578        messages: Vec<Message>,
579        mut file: Box<dyn WriteSeek + 'a>,
580        _filename: &str,
581        encoding: Encoding,
582        replacement: Option<&'a ReplacementTable>,
583    ) -> Result<()> {
584        file.write_all(b"ASB\0\0")?;
585        let mut items = self.items.clone();
586        let mut name_index = None;
587        let mut mes_index = 0;
588        let mut item_index = 0;
589        let mut print_index = None;
590        let mut hcls_index = 1;
591        while item_index < items.len() {
592            if let Some(print_ind) = print_index.clone() {
593                if items[item_index].is_command_name("hcls") {
594                    let message = messages
595                        .get(mes_index)
596                        .ok_or(anyhow::anyhow!("Not enough messages."))?;
597                    if let Some(name_index) = name_index.take() {
598                        let mut name = match &message.name {
599                            Some(name) => name.to_owned(),
600                            None => return Err(anyhow::anyhow!("Message without name.")),
601                        };
602                        if let Some(replacement) = replacement {
603                            for (k, v) in &replacement.map {
604                                name = name.replace(k, v);
605                            }
606                        }
607                        if let Item::Command(cmd) = &mut items[name_index] {
608                            if cmd.attributes.len() > 1 {
609                                cmd.attributes
610                                    .insert(format!("{}", cmd.attributes.len() - 1), name);
611                            } else {
612                                let oname = cmd
613                                    .attributes
614                                    .get("0")
615                                    .ok_or(anyhow::anyhow!("No name attribute found."))?;
616                                if oname != &name {
617                                    cmd.attributes.insert("1".to_string(), name);
618                                }
619                            }
620                        }
621                    }
622                    let mut m = message.message.clone();
623                    if let Some(replacement) = replacement {
624                        for (k, v) in &replacement.map {
625                            m = m.replace(k, v);
626                        }
627                    }
628                    let new_cmds = TextParser::new(&m.replace("\n", "<rt>"), hcls_index).parse()?;
629                    hcls_index += 1;
630                    let new_cmds_len = new_cmds.len();
631                    items.splice(print_ind..=item_index, new_cmds);
632                    print_index = None;
633                    item_index = print_ind + new_cmds_len;
634                    mes_index += 1;
635                    continue;
636                } else if items[item_index].is_command() {
637                    item_index += 1;
638                    continue;
639                }
640            }
641            if let Item::Command(cmd) = &mut items[item_index] {
642                match cmd.name.as_str() {
643                    "print" => {
644                        print_index = Some(item_index);
645                    }
646                    "name" => {
647                        name_index = Some(item_index);
648                    }
649                    "sel_text" => {
650                        let message = messages
651                            .get(mes_index)
652                            .ok_or(anyhow::anyhow!("Not enough messages."))?;
653                        let mut m = message.message.clone();
654                        if let Some(replacement) = replacement {
655                            for (k, v) in &replacement.map {
656                                m = m.replace(k, v);
657                            }
658                        }
659                        cmd.attributes.insert("text".to_string(), m);
660                        mes_index += 1;
661                    }
662                    "RegisterTextToHistory" => {
663                        let message = messages
664                            .get(mes_index)
665                            .ok_or(anyhow::anyhow!("Not enough messages."))?;
666                        let mut m = message.message.clone();
667                        if let Some(replacement) = replacement {
668                            for (k, v) in &replacement.map {
669                                m = m.replace(k, v);
670                            }
671                        }
672                        cmd.attributes.insert("1".to_string(), m);
673                        mes_index += 1;
674                    }
675                    _ => {}
676                }
677            }
678            item_index += 1;
679        }
680        if mes_index != messages.len() {
681            return Err(anyhow::anyhow!(
682                "Not all messages were processed, expected {}, got {}",
683                messages.len(),
684                mes_index
685            ));
686        }
687        file.write_u32(items.len() as u32)?;
688        for item in items {
689            file.write_item(&item, encoding)?;
690        }
691        file.flush()?;
692        Ok(())
693    }
694
695    fn custom_export(&self, filename: &std::path::Path, encoding: Encoding) -> Result<()> {
696        let s = if self.format_lua {
697            let items: Vec<_> = self
698                .items
699                .iter()
700                .map(|s| {
701                    if let Item::Command(cmd) = s {
702                        if cmd.name == "lua" {
703                            if let Some(script) = cmd.attributes.get("script") {
704                                let mut cmd = cmd.clone();
705                                cmd.attributes.insert(
706                                    "script".to_string(),
707                                    match self.format_lua(script) {
708                                        Ok(s) => s,
709                                        Err(_) => {
710                                            eprintln!("Warning: Failed to format Lua script.");
711                                            crate::COUNTER.inc_warning();
712                                            script.clone()
713                                        }
714                                    },
715                                );
716                                return Item::Command(cmd);
717                            }
718                        }
719                    }
720                    s.clone()
721                })
722                .collect();
723            self.to_string(&items)?
724        } else {
725            self.to_string(&self.items)?
726        };
727        let s = encode_string(encoding, &s, false)?;
728        let mut file = std::fs::File::create(filename)?;
729        file.write_all(&s)?;
730        Ok(())
731    }
732
733    fn custom_import<'a>(
734        &'a self,
735        custom_filename: &'a str,
736        file: Box<dyn WriteSeek + 'a>,
737        encoding: Encoding,
738        output_encoding: Encoding,
739    ) -> Result<()> {
740        create_file(
741            custom_filename,
742            file,
743            encoding,
744            output_encoding,
745            self.custom_yaml,
746        )
747    }
748}
749
750/// Creates a new ASB file.
751///
752/// * `custom_filename` - The path ot the input file.
753/// * `writer` - The writer to write the ASB script.
754/// * `encoding` - The encoding used for the ASB script.
755/// * `output_encoding` - The encoding used for the input file.
756/// * `yaml` - Whether to use YAML format instead of JSON for the input file.
757pub fn create_file<'a>(
758    custom_filename: &'a str,
759    mut writer: Box<dyn WriteSeek + 'a>,
760    encoding: Encoding,
761    output_encoding: Encoding,
762    yaml: bool,
763) -> Result<()> {
764    let f = crate::utils::files::read_file(custom_filename)?;
765    let s = decode_to_string(output_encoding, &f, true)?;
766    let items: Vec<Item> = if yaml {
767        serde_yaml_ng::from_str(&s)?
768    } else {
769        serde_json::from_str(&s)?
770    };
771    writer.write_all(b"ASB\0\0")?;
772    writer.write_u32(items.len() as u32)?;
773    for item in items {
774        writer.write_item(&item, encoding)?;
775    }
776    Ok(())
777}
778
779#[test]
780fn test_parse() {
781    let text = "Hello &lt; &amp; World!<tag><tags x=\"123\"><name 0=\"Ok\">Test";
782    let parser = TextParser::new(text, 1);
783    let items = parser.parse().unwrap();
784    assert_eq!(
785        items,
786        vec![
787            Item::Command(Command {
788                name: "print".to_string(),
789                line_number: 0,
790                attributes: [("data".to_string(), "Hello < & World!".to_string())].into(),
791            }),
792            Item::Command(Command {
793                name: "tag".to_string(),
794                line_number: 0,
795                attributes: BTreeMap::new(),
796            }),
797            Item::Command(Command {
798                name: "tags".to_string(),
799                line_number: 0,
800                attributes: [("x".to_string(), "123".to_string())].into(),
801            }),
802            Item::Command(Command {
803                name: "name".to_string(),
804                line_number: 0,
805                attributes: [("0".to_string(), "Ok".to_string())].into(),
806            }),
807            Item::Command(Command {
808                name: "print".to_string(),
809                line_number: 0,
810                attributes: [("data".to_string(), "Test".to_string())].into(),
811            }),
812            Item::Command(Command {
813                name: "hcls".to_string(),
814                line_number: 0,
815                attributes: BTreeMap::from([("0".to_string(), "1".to_string())]),
816            }),
817        ]
818    )
819}
msg_tool\scripts\artemis/asb.rs

msg_tool\scripts\artemis/
asb.rs
